import socket
from random import choice
from time import sleep
from copy import deepcopy

import sys
import math
import random
import numpy as np

from Board import Board
import Color

class MctsAgent():
    """This class describes the default Hex agent. It will randomly send a
    valid move at each turn, and it will choose to swap with a 50% chance.

    agent only maintain the state of board, the decision will made by mcts
    """

    HOST = "127.0.0.1"
    PORT = 1234

    def run(self):
        """A finite-state machine that cycles through waiting for input
        and sending moves.
        """

        self._board_size = 0
        self._board = [] # record current state of board
        self._colour = ""
        self._turn_count = 1
        self._mcts = Mcts()

        states = {
            1: MctsAgent._connect,
            2: MctsAgent._wait_start,
            3: MctsAgent._make_move,
            4: MctsAgent._wait_message,
            5: MctsAgent._close
        }

        res = states[1](self)
        while (res != 0):
            res = states[res](self)

    def _connect(self):
        """Connects to the socket and jumps to waiting for the start
        message.
        """

        self._s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._s.connect((MctsAgent.HOST, MctsAgent.PORT))

        return 2

    def _wait_start(self):
        """Initialises itself when receiving the start message, then
        answers if it is Red or waits if it is Blue.
        """

        data = self._s.recv(1024).decode("utf-8").strip().split(";")
        if (data[0] == "START"):
            self._board_size = int(data[1])
            self._board = Board(self._board_size)
            self._colour = data[2]

            if (self._colour == "R"):
                return 3
            else:
                return 4

        else:
            print("ERROR: No START message received.")
            return 0

    def _make_move(self):
        """Makes a random valid move. It will choose to swap with
        a coinflip.
        """

        if (self._turn_count == 2 and choice([0, 1]) == 1):
            msg = "SWAP\n"
        else:
            move = self._mcts.run(self._board)
            msg = f"{move[0]},{move[1]}\n"

        self._s.sendall(bytes(msg, "utf-8"))

        return 4

    def _wait_message(self):
        """Waits for a new change message when it is not its turn."""

        self._turn_count += 1

        data = self._s.recv(1024).decode("utf-8").strip().split(";")
        if (data[0] == "END" or data[-1] == "END"):
            return 5
        else:
            # CHANGE;MOVE;BOARD;COLOUR

            if (data[1] == "SWAP"):
                self._colour = self.opp_colour()
            else:
                x, y = data[1].split(",")
                self._board.set_tile_color(x, y, Color.from_char(data[-1]))

            if (data[-1] == self._colour):
                return 3

        return 4

    def _close(self):
        """Closes the socket."""

        self._s.close()
        return 0

    def opp_colour(self):
        """Returns the char representation of the colour opposite to the
        current one.
        """

        if self._colour == "R":
            return "B"
        elif self._colour == "B":
            return "R"
        else:
            return "None"


class Node:
    def __init__(self):
        pass

    @staticmethod
    def from_board(board):
        pass

class Mcts:
    def __init__(self):
        self.color = None
        self.node_dict = {}

    def board_to_node(self, board):
        """find corresponding node for current board
        """
        board_str = board.print_board()
        node = node_dict.get(board_str)
        if not node:
            node = Node.from_board(board)
        return node

    def best_move(self, node):
        return (0, 0)

    def tree_policy(self, node):
        pass

    def default_policy(self, node):
        pass

    def backup(self, node, reward):
        pass

    def run(self, board, color):
        """apply monte carlo tree search

        Parameters
            board(Board): current state of board
            color(char): color to move

        Returns
            (int, int): move
        """

        self.color = color

        # turn board into corresponding node(maybe existing)
        node = self.board_to_node(board)

        computation_budget = 10

        for _ in range(computation_budget):
            expand_node = self.tree_policy(node)
            reward = self.default_policy(expand_node)
            self.backup(expand_node, reward)

        return self.best_move(node)


if (__name__ == "__main__"):
    agent = MctsAgent()
    agent.run()
